# %%
import math

from scipy.integrate import solve_ivp

# https://docs.scipy.org/doc/scipy/reference/generated/scipy.integrate.solve_ivp.html#scipy.integrate.solve_ivp


G = 9.81  # ускорение свободного падения, м/с^2


def upward_cannon(t, y):
    """
    Функция вычисляет dy/dt, для пушки стреляющей вверх, где y - вектор состояния.

    y = [z_coord, velocity_z], т.е. координата и скорость по вертикальной оси

    dy/dt = [velocity_z, acceleration_z], т.е. скорость и ускорение по той же оси
    """

    dy_dt = [y[1], -G]  # скорость, ускорение, ось z направлена в зенит

    return dy_dt


# y0=[0, 10] - означает 0 по высоте, 10 м/с - скорость вертикально вверх в начальный
# момент времени
sol = solve_ivp(fun=upward_cannon, t_span=[0, 20], y0=[0, 100])


# %%
z_coord, velocity_z = sol.y
print(f"Скорость {velocity_z[-1]}, высота {z_coord[-1]} в конечный момент времени")


# %%
def hit_ground(t, y):
    """Event-функция, которая выдаёт нуль только при касании земли

    y = [z_coord, velocity_z],

    тогда касание поверхности земли соответствует z_coord = 0, т.е. y[0]=0
    """
    return y[0]


hit_ground.terminal = True
hit_ground.direction = -1

sol = solve_ivp(fun=upward_cannon, t_span=[0, 30], y0=[0, 100], events=hit_ground)

z_coord, velocity_z = sol.y
t = sol.t
print(f"Касание поверхности земли произошло в {t[-1]:.2f} c")
# print(f"Скорость при этом была {velocity_z[-1]} м/с, высота {z_coord[-1]} м")
print(f"Скорость при этом была {velocity_z[-1]:.2f} м/с, высота {z_coord[-1]:.2f} м")


# %%
def three_dimensional_cannon(t, y):
    """Функция вычисляет dy/dt, где y - вектор состояния.

    Parameters
    ----------
    t : float
        время
    y : list
        [rx, ry, rz, vx, vy, vz]
          0   1   2   3   4   5

    Returns
    -------
    list
        dy/dt = [vx, vy, vz, ax, ay, az]
    """

    # скорость, ускорение, ось z направлена в зенит
    dy_dt = [y[3], y[4], y[5], 0, 0, -G]

    return dy_dt


def hit_ground_three_dimensional(t, y):
    """Event-функция, которая выдаёт нуль только при касании земли

    Parameters
    ----------
    t : float
        время
    y : list
        [rx, ry, rz, vx, vy, vz]

    Returns
    -------
    float
        rz = y[2]
    """

    return y[2]


hit_ground_three_dimensional.terminal = True
hit_ground_three_dimensional.direction = -1

r0 = [0, 0, 0]
v0 = [1, 2, 100]

sol = solve_ivp(
    fun=three_dimensional_cannon,
    t_span=[0, 30],
    y0=[*r0, *v0],
    events=hit_ground_three_dimensional,
)


# %%
rx, ry, rz, vx, vy, vz = sol.y
t = sol.t
print(f"Касание поверхности земли произошло в {t[-1]:.2f} c")
print(f"rx = {rx[-1]:.2f}, ry = {ry[-1]:.2f}, rz = {rz[-1]:.2f}")
print(f"vx = {vx[-1]:.2f}, vy = {vy[-1]:.2f}, vz = {vz[-1]:.2f}")


# %%
C_F = 0.47  # https://ru.wikipedia.org/wiki/%D0%9A%D0%BE%D1%8D%D1%84%D1%84%D0%B8%D1%86%D0%B8%D0%B5%D0%BD%D1%82_%D1%81%D0%BE%D0%BF%D1%80%D0%BE%D1%82%D0%B8%D0%B2%D0%BB%D0%B5%D0%BD%D0%B8%D1%8F_%D1%84%D0%BE%D1%80%D0%BC%D1%8B
S = math.pi * 0.1**2 / 4
RHO = 1
MASS = 1


def gravity_force():
    return [0, 0, -MASS * G]


def aerodynamic_drag_force(velocity: list):
    # https://ru.wikipedia.org/wiki/Лобовое_сопротивление
    vx, vy, vz = velocity
    v_norm = math.sqrt(vx**2 + vy**2 + vz**2)

    force_norm = C_F * (RHO * v_norm**2) / 2 * S

    fx = -force_norm * vx / v_norm
    fy = -force_norm * vy / v_norm
    fz = -force_norm * vz / v_norm

    return [fx, fy, fz]


def total_force(velocity):
    grav_x, grav_y, grav_z = gravity_force()
    aero_x, aero_y, aero_z = aerodynamic_drag_force(velocity)

    force = [grav_x + aero_x, grav_y + aero_y, grav_z + aero_z]

    return force


def three_dimensional_cannon_with_drag(t, y):
    """Функция вычисляет dy/dt, где y - вектор состояния.

    Parameters
    ----------
    t : float
        время
    y : list
        [rx, ry, rz, vx, vy, vz]

    Returns
    -------
    list
        dy/dt = [vx, vy, vz, ax, ay, az]
    """

    fx, fy, fz = total_force(velocity=y[3:])

    acceleration = [fx / MASS, fy / MASS, fz / MASS]
    # скорость, ускорение, ось z направлена в зенит
    dy_dt = [*y[3:], *acceleration]

    return dy_dt


# %%
hit_ground_three_dimensional.terminal = True
hit_ground_three_dimensional.direction = -1

r0 = [0, 0, 0]
v0 = [1, 2, 100]

sol = solve_ivp(
    fun=three_dimensional_cannon_with_drag,
    t_span=[0, 30],
    y0=[*r0, *v0],
    events=hit_ground_three_dimensional,
)


rx, ry, rz, vx, vy, vz = sol.y
t = sol.t
print(f"Касание поверхности земли произошло в {t[-1]} c")
print(f"rx = {rx[-1]:.2f}, ry = {ry[-1]:.2f}, rz = {rz[-1]:.2f}")
print(f"vx = {vx[-1]:.2f}, vy = {vy[-1]:.2f}, vz = {vz[-1]:.2f}")
